# -*- coding: utf-8 -*-
# @Time    : 2018/3/31 19:20
# @Author  : Tianchiyue
# @File    : rnn.py
# @Software: PyCharm Community Edition
from models.model import BaseModel
from keras.layers import Bidirectional, LSTM, GRU


class Rnn(BaseModel):
    def build(self):
        if self.config['bidirectional']:
            if self.config['rnn'] == 'gru':
                rnn_out = Bidirectional(GRU(self.config['rnn_output_size'],
                                            dropout=self.config['dropout_rate'],
                                            recurrent_dropout=self.config['dropout_rate']))(self.sentence)
            else:
                rnn_out = Bidirectional(LSTM(self.config['rnn_output_size'],
                                             dropout=self.config['dropout_rate'],
                                             recurrent_dropout=self.config['dropout_rate']))(self.sentence)
        else:
            if self.config['rnn'] == 'gru':
                rnn_out = GRU(self.config['rnn_output_size'],
                              dropout=self.config['dropout_rate'],
                              recurrent_dropout=self.config['dropout_rate'])(self.sentence)
            else:
                rnn_out = LSTM(self.config['rnn_output_size'],
                               dropout=self.config['dropout_rate'],
                               recurrent_dropout=self.config['dropout_rate'])(self.sentence)
        return rnn_out